use std::ops::{Deref, DerefMut, Index, IndexMut};

use mark_dirty::DirtyMark;

use crate::{math::RegionRect, scene::AsBytes};

use super::Attribute;

#[repr(C)]
#[derive(Clone, Copy)]
pub struct Pixel {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

impl Pixel {
    pub const TRANSPARENT: Self = Self::new(0, 0, 0, 0);
    pub const BLACK: Self = Self::new(0, 0, 0, 255);
    pub const BLUE: Self = Self::new(0, 0, 255, 255);
    pub const GREEN: Self = Self::new(0, 255, 0, 255);
    pub const CYAN: Self = Self::new(0, 255, 255, 255);
    pub const RED: Self = Self::new(255, 0, 0, 255);
    pub const PURPLE: Self = Self::new(255, 0, 255, 255);
    pub const YELLOW: Self = Self::new(255, 255, 0, 255);
    pub const WHITE: Self = Self::new(255, 255, 255, 255);

    pub const fn new(r: u8, g: u8, b: u8, a: u8) -> Self {
        Self { r, g, b, a }
    }

    pub const fn from_rgba_array(rgba: [u8; 4]) -> Self {
        Self {
            r: rgba[0],
            g: rgba[1],
            b: rgba[2],
            a: rgba[3],
        }
    }
}

impl AsBytes for Pixel {}
// TODO: Which endian the computer have? if we could know, simply use unsafe { transmute(value) }
impl From<u32> for Pixel {
    fn from(value: u32) -> Self {
        let mask = 0b11u32;
        let a = (value & mask) as u8;
        let b = ((value >> 8) & mask) as u8;
        let g = ((value >> 16) & mask) as u8;
        let r = ((value >> 24) & mask) as u8;

        Self { r, g, b, a }
    }
}

pub struct Pixmap {
    pub pixels: Vec<Pixel>,
    pub width: u32,
    pub height: u32,
}

impl Pixmap {
    pub fn new(width: u32, height: u32) -> Self {
        Self {
            pixels: vec![Pixel::TRANSPARENT; (width * height) as usize],
            width,
            height,
        }
    }

    pub fn from_pixels<I: IntoIterator<Item = Pixel>>(pixels: I, width: u32, height: u32) -> Self {
        Self {
            pixels: pixels.into_iter().collect(),
            width,
            height,
        }
    }

    pub fn from_pixel_array_2d(pixels: &[&[Pixel]]) -> Self {
        let height = pixels.len() as u32;
        let width = pixels[0].len() as u32;
        let pixels = pixels
            .iter()
            .map(|l| l.iter())
            .flatten()
            .map(|p| *p)
            .collect::<Vec<_>>();

        Self {
            pixels,
            width,
            height,
        }
    }

    pub fn create_region(&self, x: u32, y: u32, width: u32, height: u32) -> RegionRect {
        let pixmap_width = self.width as f32;
        let pixmap_height = self.height as f32;
        let left = x as f32 / pixmap_width;
        let top = y as f32 / pixmap_height;
        let right = (x + width) as f32 / pixmap_width;
        let bottom = (y + height) as f32 / pixmap_height;

        RegionRect::new(left, bottom, right, top)
    }

    pub fn extent(&self) -> wgpu::Extent3d {
        wgpu::Extent3d {
            width: self.width,
            height: self.height,
            depth_or_array_layers: 1,
        }
    }

    pub fn data_layout(&self) -> wgpu::ImageDataLayout {
        wgpu::ImageDataLayout {
            offset: 0,
            bytes_per_row: Some(std::mem::size_of::<Pixel>() as u32 * self.width),
            rows_per_image: Some(self.height),
        }
    }
}

impl Index<(u32, u32)> for Pixmap {
    type Output = Pixel;

    fn index(&self, index: (u32, u32)) -> &Self::Output {
        &self.pixels[(index.0 + self.width * index.1) as usize]
    }
}

impl IndexMut<(u32, u32)> for Pixmap {
    fn index_mut(&mut self, index: (u32, u32)) -> &mut Self::Output {
        &mut self.pixels[(index.0 + self.width * index.1) as usize]
    }
}

pub struct Texture {
    pixmap: DirtyMark<Pixmap>,
    layout: wgpu::BindGroupLayout,
    bind_group: wgpu::BindGroup,
    texture: wgpu::Texture,
}

impl Texture {
    pub fn new(ctx: &super::Context, visibility: wgpu::ShaderStages, pixmap: Pixmap) -> Self {
        let sampler = ctx.device.create_sampler(&wgpu::SamplerDescriptor {
            min_filter: wgpu::FilterMode::Nearest,
            mag_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });

        Self::with_sampler(ctx, pixmap, &sampler, visibility)
    }

    pub fn with_sampler(
        ctx: &super::Context,
        pixmap: Pixmap,
        sampler: &wgpu::Sampler,
        visibility: wgpu::ShaderStages,
    ) -> Self {
        let layout = ctx
            .device
            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: None,
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility,
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                        count: None,
                    },
                ],
            });
        let texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
            label: None,
            size: pixmap.extent(),
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8UnormSrgb,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });

        let bind_group = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(
                        &texture.create_view(&wgpu::TextureViewDescriptor::default()),
                    ),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
            ],
        });

        let mut pixmap = DirtyMark::new(pixmap);
        pixmap.mark_dirty();

        Self {
            pixmap,
            layout,
            bind_group,
            texture,
        }
    }
}

impl Attribute for Texture {
    fn layout(&self) -> &wgpu::BindGroupLayout {
        &self.layout
    }

    fn bind_group(&self) -> &wgpu::BindGroup {
        &self.bind_group
    }

    fn prepare(&mut self, ctx: &super::Context) {
        if self.pixmap.is_dirty() {
            ctx.queue.write_texture(
                wgpu::ImageCopyTextureBase {
                    texture: &self.texture,
                    mip_level: 0,
                    origin: wgpu::Origin3d::ZERO,
                    aspect: wgpu::TextureAspect::All,
                },
                self.pixmap.pixels.as_slice().as_bytes(),
                self.pixmap.data_layout(),
                self.texture.size(),
            );
            self.pixmap.unmark_dirty()
        }
    }
}

impl Deref for Texture {
    type Target = Pixmap;

    fn deref(&self) -> &Self::Target {
        &self.pixmap
    }
}

impl DerefMut for Texture {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.pixmap
    }
}
